from src.model import APIModel
import re
import json
import time
from tqdm import tqdm

def generate_topics(output_file="src/data/concept_topics.txt", batch_size=50, total=2500):
    model = APIModel(model="glm-4-flash")
    base_prompt1 = """
Generate technical concept terms in the field of computer science. Requirements:
1. Focus on specific technical points (e.g., adversarial examples, consensus algorithms). Only output the technical concept term, without any additional content (e.g., "applications of xxx," "trends of xxx").
2. Cover diverse research areas (e.g., machine learning, systems, security).
3. Follow the format example: "Loss Functions"

Output the terms directly, without numbering, with each term on a separate line. Generate {} distinct terms.
"""

    generated = set()
    batch_count = total // batch_size

    for i in tqdm(range(batch_count)):
        retries = 3
        while retries > 0:
            try:
                prompt = base_prompt1.format(batch_size)
                response = model.chat(prompt)

                # 后处理步骤
                new_topics = [re.sub(r'^\d+[\.\)]\s*', '', t).strip() for t in response.split("\n")]

                new_topics = list(set(new_topics))

                # 去重处理
                unique_new = [t for t in new_topics if t not in generated]

                if len(unique_new) >= batch_size * 0.8:  # 允许部分重复
                    with open(output_file, "a", encoding="utf-8") as f:
                        f.write("\n".join(unique_new) + "\n")
                    generated.update(unique_new)
                    print(f"批次 {i + 1}/{batch_count} 完成，新增 {len(unique_new)} 条")
                    break
                else:
                    retries -= 1
                    print(f"批次 {i + 1} 重复率过高，剩余重试次数 {retries}")
            except Exception as e:
                print(f"接口调用失败: {str(e)}")
                retries -= 1
                time.sleep(5)

        #time.sleep(1.5)  # 控制请求频率

    print(f"生成完成，总计生成 {len(generated)} 个唯一主题")


# 执行生成
generate_topics()

